import gradio as gr
import torch
from torchvision import transforms
from PIL import Image
import torchvision.models as models
import torch.nn as nn
import MyNet
# 定义数据预处理的转换
mean = [0.5, 0.5, 0.5]
# std = [0.25, 0.25, 0.25]
std = [0.5, 0.5, 0.5]
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean, std),
])
label_Butterfly = {
    0:'ADONIS',
    1:'AFRICAN GIANT SWALLOWTAIL',
    2:'AMERICAN SNOOT',
    3:'AN 88',
    4:'APPOLLO',
    5:'ATALA',
    6:'BANDED ORANGE HELICONIAN',
    7:'BANDED PEACOCK',
    8:'BECKERS WHITE',
    9:'BLACK HAIRSTREAK',
    10:'BLUE MORPHO',
    11:'BLUE SPOTTED CROW',
    12:'BROWN SIPROETA',
    13:'CABBAGE WHITE',
    14:'CAIRNS BIRDWING',
    15:'CHECQUERED SKIPPER',
    16:'CHESTNUT',
    17:'CLEOPATRA',
    18:'CLODIUS PARNASSIAN',
    19:'CLOUDED SULPHUR',
    20:'COMMON BANDED AWL',
    21:'COMMON WOOD-NYMPH',
    22:'COPPER TAIL',
    23:'CRECENT',
    24:'CRIMSON PATCH',
    25:'DANAID EGGFLY',
    26:'EASTERN COMA',
    27:'EASTERN DAPPLE WHITE',
    28:'EASTERN PINE ELFIN',
    29:'ELBOWED PIERROT',
    30:'GOLD BANDED',
    31:'GREAT EGGFLY',
    32:'GREAT JAY',
    33:'GREEN CELLED CATTLEHEART',
    34:'GREY HAIRSTREAK',
    35:'INDRA SWALLOW',
    36:'IPHICLUS SISTER',
    37:'JULIA',
    38:'LARGE MARBLE',
    39:'MALACHITE',
    40:'MANGROVE SKIPPER',
    41:'MESTRA',
    42:'METALMARK',
    43:'MILBERTS TORTOISESHELL',
    44:'MONARCH',
    45:'MOURNING CLOAK',
    46:'ORANGE OAKLEAF',
    47:'ORANGE TIP',
    48:'ORCHARD SWALLOW',
    49:'PAINTED LADY',
    50:'PAPER KITE',
    51:'PEACOCK',
    52:'PINE WHITE',
    53:'PIPEVINE SWALLOW',
    54:'POPINJAY',
    55:'PURPLE HAIRSTREAK',
    56:'PURPLISH COPPER',
    57:'QUESTION MARK',
    58:'RED ADMIRAL',
    59:'RED CRACKER',
    60:'RED POSTMAN',
    61:'RED SPOTTED PURPLE',
    62:'SCARCE SWALLOW',
    63:'SILVER SPOT SKIPPER',
    64:'SLEEPY ORANGE',
    65:'SOOTYWING',
    66:'SOUTHERN DOGFACE',
    67:'STRAITED QUEEN',
    68:'TROPICAL LEAFWING',
    69:'TWO BARRED FLASHER',
    70:'ULYSES',
    71:'VICEROY',
    72:'WOOD SATYR',
    73:'YELLOW SWALLOW TAIL',
    74:'ZEBRA LONG WING'
}
num_classes = 75
device1 = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
print(device1)
parameter = '_'+str(30)+'_'+str(32)+'_'+str(0.0005)+'_'+str(0.9)
def get_alexnet():
    model_alexnet = models.alexnet(weights=None)
    model_alexnet.classifier[6] = nn.Linear(model_alexnet.classifier[6].in_features, num_classes)
    model_alexnet.load_state_dict(torch.load('models/butterfly_AlexNet'+parameter+'.pth', map_location=device1))
    model_alexnet.eval()
    return model_alexnet
def get_VGG():
    model = models.vgg19(weights=None)
    model.classifier[6] = nn.Linear(model.classifier[6].in_features, num_classes)
    model.load_state_dict(torch.load('models/butterfly_VGG19'+parameter+'.pth', map_location=device1))
    model.eval()
    return model
def get_ResNet50():
    model = models.resnet50(weights=None)
    model.fc = nn.Linear(model.fc.in_features, num_classes)
    model.load_state_dict(torch.load('models/butterfly_ResNet50'+parameter+'.pth', map_location=device1))
    model.eval()
    return model
def get_DenseNet121():
    model = models.densenet121(weights=None)
    num_features = model.classifier.in_features
    model.classifier = nn.Linear(num_features, num_classes)
    model.load_state_dict(torch.load('models/butterfly_DenseNet121'+parameter+'.pth', map_location=device1))
    model.eval()
    return model

def get_SqueezeNet():
    model = models.squeezenet1_1(weights=None)
    # 获取模型最后一个层的输入特征数
    num_features = model.classifier[1].in_channels
    # 替换最后一个全连接层
    model.classifier[1] = nn.Conv2d(num_features, num_classes, kernel_size=(1, 1), stride=(1, 1))
    # 将模型的最后一层设置为分类层
    model.num_classes = num_classes
    model.load_state_dict(torch.load('models/butterfly_SqueezeNet'+parameter+'.pth', map_location=device1))
    model.eval()
    return model

def get_MyNet():
    model = MyNet.MyModel1()
    model.load_state_dict(torch.load('models/butterfly_MyNet'+parameter+'.pth', map_location=device1))
    model.eval()
    return model

model_list =  ["AlexNet", "VGG19", "ResNet50", "DenseNet121", "SqueezeNet", "MyNet"]
model_funs = {
        "AlexNet": get_alexnet,
        "VGG19": get_VGG,
        "ResNet50": get_ResNet50,
        "DenseNet121": get_DenseNet121,
        "SqueezeNet": get_SqueezeNet,
        "MyNet": get_MyNet
    }
def predictproocess(model,img):
    img = transform(img)  # 应用数据预处理的转换
    # 进行预测
    with torch.no_grad():
        output = model(img.unsqueeze(0))  # 将图像增加一个维度作为批处理维度
        predicted_class = torch.argmax(output).item()
    return predicted_class

def greet(name,image):
    models = model_list
    if name in models:
        model = change_model(name)
        image = Image.fromarray(image.astype('uint8'), 'RGB')
        predicted_class = predictproocess(model, image)
        return label_Butterfly[predicted_class]


def change_model(model_name):
    models = model_list
    model_functions = model_funs
    if model_name in models:
        # 在模型列表中找到匹配项，执行相应函数或操作
        model = model_functions[model_name]()
        return model
    else:
        print("Name not found in models list")


dropdown = gr.Dropdown(model_list, label="Models", info="Change model" )
demo = gr.Interface(fn=greet,
                    inputs=[dropdown,"image"],
                    outputs="text",
                    title="Image Processing Interface",
                   )
if __name__ == "__main__":
    demo.launch(show_api=False)
    print("ok")